import os
import cv2
import numpy as np
import gradio as gr
import logging
from PIL import Image

from database import FaceDatabase
from models import SCRFD, ArcFace
from utils.logging import setup_logging
from utils.helpers import compute_similarity, draw_bbox_info, draw_bbox

setup_logging(log_to_file=True)

class FaceIdentificationApp:
    def __init__(self):
        self.det_weight = "./weights/det_10g.onnx"
        self.rec_weight = "./weights/w600k_r50.onnx"
        self.similarity_thresh = 0.4
        self.confidence_thresh = 0.5
        self.faces_dir = "./assets/faces"
        self.db_path = "./database/face_database"
        
        self.detector = None
        self.recognizer = None
        self.face_db = None
        self.colors = {}
        
        self._initialize_models()
        
    def _initialize_models(self):
        try:
            logging.info("Initializing face detection and recognition models...")
            self.detector = SCRFD(self.det_weight, input_size=(640, 640), conf_thres=self.confidence_thresh)
            self.recognizer = ArcFace(self.rec_weight)
            logging.info("Models initialized successfully")
            
            self._build_face_database()
            
        except Exception as e:
            logging.error(f"Failed to initialize models: {e}")
            raise
    
    def _build_face_database(self):
        self.face_db = FaceDatabase(db_path=self.db_path)
        
        if self.face_db.load():
            logging.info("Loaded existing face database from disk.")
            return
        
        logging.info("Building face database from images...")
        if not os.path.exists(self.faces_dir):
            logging.warning(f"Faces directory {self.faces_dir} does not exist")
            return
            
        for filename in os.listdir(self.faces_dir):
            if not (filename.endswith('.jpg') or filename.endswith('.png')):
                continue
                
            name = filename.rsplit('.', 1)[0]
            image_path = os.path.join(self.faces_dir, filename)
            image = cv2.imread(image_path)
            
            if image is None:
                logging.warning(f"Could not read image: {image_path}")
                continue
                
            bboxes, kpss = self.detector.detect(image, max_num=1)
            
            if len(kpss) == 0:
                logging.warning(f"No face detected in {image_path}. Skipping...")
                continue
                
            embedding = self.recognizer.get_embedding(image, kpss[0])
            self.face_db.add_face(embedding, name)
            logging.info(f"Added face for: {name}")
        
        self.face_db.save()
    
    def identify_face(self, image):
        if image is None:
            return None, "Please upload an image"
        
        try:
            if isinstance(image, Image.Image):
                image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            
            bboxes, kpss = self.detector.detect(image, max_num=0)
            
            if len(kpss) == 0:
                return image, "No faces detected in the image"
            
            results = []
            annotated_image = image.copy()
            
            for bbox, kps in zip(bboxes, kpss):
                *bbox, conf_score = bbox.astype(np.int32)
                embedding = self.recognizer.get_embedding(image, kps)
                
                name, similarity = self.face_db.search(embedding, self.similarity_thresh)
                
                if name != "Unknown":
                    if name not in self.colors:
                        import random
                        self.colors[name] = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
                    draw_bbox_info(annotated_image, bbox, similarity=similarity, name=name, color=self.colors[name])
                    results.append(f"Identified: {name} (Similarity: {similarity:.3f})")
                else:
                    draw_bbox(annotated_image, bbox, (255, 0, 0))
                    results.append(f"Unknown person (Confidence: {conf_score:.3f})")
            
            result_text = "\n".join(results) if results else "No faces identified"
            
            annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
            return annotated_image_rgb, result_text
            
        except Exception as e:
            logging.error(f"Error during face identification: {e}")
            return image, f"Error during identification: {str(e)}"
    
    def add_new_face(self, image, name):
        if image is None:
            return "Please upload an image"
        
        if not name or name.strip() == "":
            return "Please enter a valid name"
        
        name = name.strip()
        
        try:
            if isinstance(image, Image.Image):
                image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            
            bboxes, kpss = self.detector.detect(image, max_num=1)
            
            if len(kpss) == 0:
                return "No face detected in the image. Please upload a clear image with a visible face."
            
            if len(kpss) > 1:
                logging.warning("Multiple faces detected, using the first one")
            
            embedding = self.recognizer.get_embedding(image, kpss[0])
            
            existing_name, similarity = self.face_db.search(embedding, self.similarity_thresh)
            if existing_name != "Unknown":
                return f"This face is already registered as '{existing_name}' (similarity: {similarity:.3f})"
            
            self.face_db.add_face(embedding, name)
            self.face_db.save()
            
            os.makedirs(self.faces_dir, exist_ok=True)
            image_filename = f"{name}.png"
            image_path = os.path.join(self.faces_dir, image_filename)
            cv2.imwrite(image_path, image)
            
            logging.info(f"Added new face for: {name}")
            return f"Successfully added '{name}' to the database!"
            
        except Exception as e:
            logging.error(f"Error adding new face: {e}")
            return f"Error adding face: {str(e)}"

def create_interface():
    app = FaceIdentificationApp()
    
    with gr.Blocks(title="Face Re-identification System") as interface:
        gr.Markdown("# Face Re-identification System")
        gr.Markdown("Upload an image to identify faces or add new faces to the database")
        
        with gr.Tabs():
            with gr.TabItem("Face Identification"):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### Upload Image for Identification")
                        input_image = gr.Image(
                            type="pil",
                            label="Upload an image containing faces"
                        )
                        identify_btn = gr.Button("Identify Faces", variant="primary")
                    
                    with gr.Column():
                        gr.Markdown("### Identification Results")
                        output_image = gr.Image(label="Annotated Image")
                        result_text = gr.Textbox(
                            label="Identification Results",
                            lines=5,
                            interactive=False
                        )
                
                identify_btn.click(
                    fn=app.identify_face,
                    inputs=[input_image],
                    outputs=[output_image, result_text]
                )
            
            with gr.TabItem("Add New Person"):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### Add New Person to Database")
                        new_face_image = gr.Image(
                            type="pil",
                            label="Upload a clear image of the person's face"
                        )
                        person_name = gr.Textbox(
                            label="Person's Name",
                            placeholder="Enter the person's name"
                        )
                        add_btn = gr.Button("Add to Database", variant="primary")
                    
                    with gr.Column():
                        gr.Markdown("### Add Person Results")
                        add_result = gr.Textbox(
                            label="Status",
                            lines=3,
                            interactive=False
                        )
                
                add_btn.click(
                    fn=app.add_new_face,
                    inputs=[new_face_image, person_name],
                    outputs=[add_result]
                )
        
        gr.Markdown("### Instructions:")
        gr.Markdown("**Face Identification:**")
        gr.Markdown("1. Upload an image containing one or more faces")
        gr.Markdown("2. Click 'Identify Faces' to run the identification")
        gr.Markdown("3. View the annotated image and identification results")
        gr.Markdown("4. Known faces will be highlighted with colored bounding boxes and names")
        gr.Markdown("5. Unknown faces will be highlighted with red bounding boxes")
        gr.Markdown("")
        gr.Markdown("**Add New Person:**")
        gr.Markdown("1. Upload a clear image with a single face")
        gr.Markdown("2. Enter the person's name")
        gr.Markdown("3. Click 'Add to Database' to register the new person")
        gr.Markdown("4. The person will be available for future identifications")
    
    return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch()